from pymongo import MongoClient
from typing import Any, Dict, List, Union
from utils.config import config


class MongoDBClient:
    def __init__(self, uri: str = config.mongo_uri, db_name: str = config.mongo_db_name):
        """
        Initialize the MongoDB client.

        :param uri: MongoDB connection URI. `[Default: 'config.mongo_uri']`
        :param db_name: Database name. `[Default: 'config.mongo_db_name']`
        """
        self.client = self.connect(uri)
        self.db = self.client[db_name]

    def connect(self, uri: str) -> MongoClient:
        """
        Connect to the MongoDB server using a connection URI.

        :param uri: MongoDB connection URI.
        :return: Connected MongoClient object.
        """
        try:
            client = MongoClient(uri)
            # 验证连接是否成功
            client.admin.command('ping')
            print("Connected to MongoDB successfully.")
            return client
        except Exception as e:
            print(f"Failed to connect to MongoDB: {e}")
            raise

    def insert_one(self, collection_name: str, document: Dict[str, Any]) -> None:
        """
        Insert a single document into a collection.

        :param collection_name: Collection name.
        :param document: Document to insert.
        """
        collection = self.db[collection_name]
        result = collection.insert_one(document)
        print(f"Inserted document with ID: {result.inserted_id}")

    def insert_many(self, collection_name: str, documents: List[Dict[str, Any]]) -> None:
        """
        Insert multiple documents into a collection.

        :param collection_name: Collection name.
        :param documents: List of documents to insert.
        """
        collection = self.db[collection_name]
        result = collection.insert_many(documents)
        print(f"Inserted {len(result.inserted_ids)} documents")

    def find_one(self, collection_name: str, query: Dict[str, Any] = {}) -> Dict[str, Any]:
        """
        Find a single document matching the query.

        :param collection_name: Collection name.
        :param query: Query dictionary.
        :return: Matching document or None if no match is found.
        """
        collection = self.db[collection_name]
        return collection.find_one(query)

    def find(self, collection_name: str, query: Dict[str, Any] = {}) -> List[Dict[str, Any]]:
        """
        Find all documents matching the query.

        :param collection_name: Collection name.
        :param query: Query dictionary.
        :return: List of matching documents.
        """
        collection = self.db[collection_name]
        return list(collection.find(query))

    def update_one(self, collection_name: str, query: Dict[str, Any], new_values: Dict[str, Any]):
        """
        Update a single document matching the query.

        :param collection_name: Collection name.
        :param query: Query dictionary to match the document.
        :param new_values: Dictionary with new values.
        """
        collection = self.db[collection_name]
        result = collection.update_one(query, {"$set": new_values})
        print(f"{result.modified_count} document(s) updated")

    def update_many(self, collection_name: str, query: Dict[str, Any], new_values: Dict[str, Any]):
        """
        Update multiple documents matching the query.

        :param collection_name: Collection name.
        :param query: Query dictionary to match the document.
        :param new_values: Dictionary with new values.
        """
        collection = self.db[collection_name]
        result = collection.update_many(query, {"$set": new_values})
        print(f"{result.modified_count} document(s) updated")

    def delete_one(self, collection_name: str, query: Dict[str, Any]):
        """
        Delete a single document matching the query.

        :param collection_name: Collection name.
        :param query: Query dictionary to match the document.
        """
        collection = self.db[collection_name]
        result = collection.delete_one(query)
        print(f"{result.deleted_count} document(s) deleted")

    def delete_many(self, collection_name: str, query: Dict[str, Any]):
        """
        Delete multiple documents matching the query.

        :param collection_name: Collection name.
        :param query: Query dictionary to match the document.
        """
        collection = self.db[collection_name]
        result = collection.delete_many(query)
        print(f"{result.deleted_count} document(s) deleted")


# 示例用法
if __name__ == "__main__":
    client = MongoDBClient()

    data = [
        {"username": "root", "password": 25},
    ]

    client.insert_many("users", data)

    doc = client.find_one("users", {"name": "Alice"})
    print(doc)

    doc["age"] = 26
    client.update_one("users", {"name": "Alice"}, doc)

    updated_doc = client.find_one("users", {"name": "Alice"})
    print(updated_doc)

    client.delete_one("users", {"name": "Bob"})

    remaining_docs = client.find("users")
    for doc in remaining_docs:
        print(doc)
